# -*- coding: utf-8 -*-
# created on 2016/12/12

from mathsolver.functions.base import *
from mathsolver.functions.shulie.Distinguish import *
from mathsolver.functions.shulie.ShuLie_getSums import ShuLieGetSums


class ShuLieGetMuls001(BaseFunction):
    """
    求an=(an+b)/(an+c)的乘积
    """
    def solver(self, *args):
        poly = args[0]
        num = args[1]
        symbol = sl_default_symbol(poly)
        judge, pcoeff, fenzhi, fenmu = sl_liancheng001(poly, symbol)
        assert judge
        if fenzhi - fenmu > 0:
            jianju = fenzhi - fenmu
            part1 = 1
            part2 = 1
            for i in range(1, jianju + 1):
                part1 = part1 * fenmu.subs({symbol: i})
                part2 = part2 * fenzhi.subs({symbol: num - 1 + i})
            mul_n = part2 / part1 * pcoeff ** num
        else:
            jianju = fenmu - fenzhi
            part1 = 1
            part2 = 1
            for i in range(1, jianju + 1):
                part2 = part2 * fenzhi.subs({symbol: i})
                part1 = part1 * fenmu.subs({symbol: num - 1 + i})
            mul_n = part2 / part1 * pcoeff ** num
        return mul_n


class ShuLieGetMuls002(BaseFunction):
    """
    求an=p*a**r的乘积
    """
    def solver(self, *args):
        poly = args[0]
        num = args[1]
        symbol = sl_default_symbol(poly)
        judge, pcoeff, pbase, pexp = sl_liancheng002(poly, symbol)
        assert judge
        exp_sums = ShuLieGetSums().solver([sympify('a_(n)'), pexp], num).output[0].toSympy()[0]
        mul_n = pcoeff ** num * pbase ** exp_sums
        mul_n = mul_n.simplify()
        return mul_n


#
class ShuLieGetMul(BaseFunction):
    CLS = [ShuLieGetMuls001, ShuLieGetMuls002]

    def solver(self, *args):
        r = None
        for cl in ShuLieGetMul.CLS:
            try:
                r = cl(verbose=True).solver(*args)
                break
            except Exception:
                pass
        if not r:
            raise 'try fail'
        return r


if __name__ == '__main__':
    pass
